import time
import pygame
import numpy as np


class Draw(object):
    def __init__(self, matrix):
        self.white = (180, 180, 180)
        self.black = (73, 37, 24)
        self.red = (255, 0, 0)
        self.width = 500
        self.height = 500
        self.grid_num = matrix.shape[1]
        self.grid_size = self.width // self.grid_num
        self.matrix = matrix
        self.robot_li = np.argwhere(matrix[0] == 1)
        self.start_li = np.argwhere(matrix[1] == 1)
        self.end_li = np.argwhere(matrix[2] == 1)
        self.stone_li = np.argwhere(matrix[3] == 1)
        pygame.init()
        self.screen = pygame.display.set_mode([self.width, self.height])
        stone_img = pygame.image.load("data/stone.png")
        self.stone_img = pygame.transform.scale(stone_img, (self.grid_size, self.grid_size))

        agent_img = pygame.image.load("data/agent.png")
        self.agent_img = pygame.transform.scale(agent_img, (self.grid_size, self.grid_size))

        one_img = pygame.image.load("data/1.png")
        self.one_img = pygame.transform.scale(one_img, (self.grid_size, self.grid_size))

        two_img = pygame.image.load("data/2.png")
        self.two_img = pygame.transform.scale(two_img, (self.grid_size, self.grid_size))

        three_img = pygame.image.load("data/3.png")
        self.three_img = pygame.transform.scale(three_img, (self.grid_size, self.grid_size))

        four_img = pygame.image.load("data/4.png")
        self.four_img = pygame.transform.scale(four_img, (self.grid_size, self.grid_size))

        five_img = pygame.image.load("data/5.png")
        self.five_img = pygame.transform.scale(five_img, (self.grid_size, self.grid_size))

        self.update(self.matrix)

    def update(self, matrix):
        self.matrix = matrix
        self.draw_background()
        self.draw_agent_start_end()

    def draw_agent_start_end(self):
        for i in np.argwhere(self.matrix[3] == 1):
            self.screen.blit(self.stone_img, [i[0] * self.grid_size, i[1] * self.grid_size])
        for i in np.argwhere(self.matrix[0] == 1):
            self.screen.blit(self.agent_img, [i[0] * self.grid_size, i[1] * self.grid_size])
        for i in np.argwhere(self.matrix[1] == 1):
            self.screen.blit(self.one_img, [i[0] * self.grid_size, i[1] * self.grid_size])
        for i in np.argwhere(self.matrix[2] == 1):
            self.screen.blit(self.two_img, [i[0] * self.grid_size, i[1] * self.grid_size])
        for i in np.argwhere(self.matrix[4] == 1):
            self.screen.blit(self.three_img, [i[0] * self.grid_size, i[1] * self.grid_size])
        for i in np.argwhere(self.matrix[5] == 1):
            self.screen.blit(self.four_img, [i[0] * self.grid_size, i[1] * self.grid_size])
        for i in np.argwhere(self.matrix[6] == 1):
            self.screen.blit(self.five_img, [i[0] * self.grid_size, i[1] * self.grid_size])

        pygame.display.update()

    def draw_background(self):
        self.screen.fill(self.black)
        for i in range(0, self.width, self.grid_size):
            pygame.draw.line(self.screen, self.white, (0, i), (self.width, i))
            pygame.draw.line(self.screen, self.white, (i, 0), (i, self.width))
        for i in self.stone_li:
            self.screen.blit(self.stone_img, [i[0] * self.grid_size, i[1] * self.grid_size])
        pygame.display.update()


if __name__ == '__main__':
    matrix = np.zeros([4, 10, 10])
    ## 小车
    matrix[0][2, 2] = 1
    matrix[0][8, 8] = 1
    ## 乘客A
    matrix[1][0, 8] = 1
    ## 乘客B
    matrix[2][8, 1] = 1
    ## 障碍
    for i in range(4):
        for j in range(4):
            matrix[3][j+3,i] = 1
            matrix[3][j+3,i+6] = 1
    # matrix[3][2, 5] = 1
    # matrix[3][5, 6] = 1
    # matrix[3][2, 6] = 1
    # matrix[3][6, 8] = 1
    dr = Draw(matrix)
    time.sleep(30)




